from numpy import *
def qrdecom(Ain,ert=100*finfo(float).eps):
	""" A fast and stable QR decomposition. Ain is a n by m matrix,
	with n>=m. """
	A=Ain.astype('float')	
	n=A.shape[0]
	m=A.shape[1]
	if n<m:
		print "Number of rows of input matrix should be larger than that of columns!"
		return []
	if m==1:
		v=zeros(A.shape)
		v[0,0]=1
		if max(abs(A))<=ert:
			R=zeros(A.shape)
			nv=dot(v.T,v)		
			W=v/nv**0.5
			Y=2*v.T/nv**0.5
		else:
			if A[0,0]!=0:			
				R=-sign(A[0,0])*linalg.norm(A[:,0])*v
			else:
				R=-linalg.norm(A[:,0])*v
			v=A-R
			nv=dot(v.T,v)		
			W=v/nv**0.5
			Y=2*v.T/nv**0.5			
		return [R,W,Y]
	else:
		Vl=qrdecom(A[:,0:m/2])
		A[:,m/2:m]=A[:,m/2:m]-mat(Vl[1])*mat(Vl[2])*mat(A[:,m/2:m])
		Vr=qrdecom(A[m/2:n,m/2:m])
		X=Vl[1]-vstack([zeros((m/2,m/2)),mat(Vr[1])*(mat(Vr[2])*mat(Vl[1][m/2:n,0:m/2]))])
		R=vstack([hstack([Vl[0][0:m/2,:],A[0:m/2,m/2:m]]),hstack([zeros((n-m/2,m/2)),Vr[0]])])
		W=hstack([X,vstack([zeros((m/2,m-m/2)),Vr[1]])])
		Y=vstack([Vl[2],hstack([zeros((m-m/2,m/2)),Vr[2]])])
		return [R,W,Y]
		
	
def backsubup(R,b,ert=100*finfo(float).eps):
	""" Back substitution to solve linear equations with upper triangle coefficient matrix"""
	x=zeros(b.shape)
	m=b.size
	if m==1:
		if R[0,0]<ert:
			return x
		else:
			x[0]=float(b[0])/R[0,0]
			return x
	else:
		if abs(R[m-1,m-1])<ert:
			if m-1>0:
				x[0:m-1]=backsubup(R[0:m-1,0:m-1],b[0:m-1])
			else:
				return x
		else:
			x[m-1]=float(b[m-1])/R[m-1,m-1]
			bm=b[0:m-1]-multiply(x[m-1],R[0:m-1,m-1])
			if m-1>0:
				x[0:m-1]=backsubup(R[0:m-1,0:m-1],bm)
	return x

def qrlstsq(A,b,ert=100*finfo(float).eps):
	R,W,Y=qrdecom(A,ert)
	n=A.shape[0]
	m=A.shape[1]
	matb=mat(b)
	if matb.shape[0]==n:
		bm=(eye(n)-mat(W)*mat(Y))*matb
	else:
		bm=(eye(n)-mat(W)*mat(Y))*matb.T
	Ru=mat(R[0:m,:])
	bu=bm[0:m]
	x=backsubup(Ru,bu,ert)
	return x
def tiklstsq(A,b,rcond=1e-10):
	N=min(A.shape)	
	U,S,V=linalg.svd(A)
	h=rcond*S[0]
	y=dot(U.T,b).reshape(b.shape)
	for i in range(0,N):
		y[i]=y[i]*S[i]/(S[i]**2+h**2)
	y=dot(V.T,y).reshape(A.shape[1])
	return y

if __name__=="__main__":
	x=arange(0,128)
	y=x+random.normal(0,1,128)
#	b=array([[2],[6]])
#	x=backsubup(A,b)
#	print x
#	R,W,Y=qrdecom(A)
#	print "R is: ", R
#	print "W is: ", W
#	print "Y is: ", Y
	A=ones([128,2])
	A[:,1]=x
	f=qrlstsq(A,y)
	print f
